import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
import networkx as nx
import os

root = os.path.dirname(__file__)
plt.rcParams["font.family"] = "SimHei"
plt.rcParams["axes.unicode_minus"] = False


def generate_nodes(xy, wh, step=10):
    """
    生成给定矩形的10x10小正方形的中心坐标。

    参数:
    xy (tuple): 矩形的左下角坐标 (x0, y0)。
    wh (tuple): 矩形的宽度和高度 (w, h)。

    返回:
    list: 所有小正方形的中心坐标 [(x, y), ...]。
    """
    x0, y0 = xy
    w, h = wh
    assert w % step == 0 and h % step == 0
    X, Y = np.meshgrid(
        np.arange(x0 + step / 2, x0 + w, step), np.arange(y0 + step / 2, y0 + h, step)
    )
    return np.array(list(zip(X.ravel(), Y.ravel())))


people = 10000
velocity = {"youth": 1.28, "elderly": 1.03, "disable": 0.78}
percent = {"youth": 0.721, "elderly": 0.211, "disable": 0.068}
velocity_mean = (np.array(list(velocity.values())) * np.array(list(percent.values()))).sum()
alpha = 0.5
start = 0
viewpoint = 22
stairs = {
    "length": (length := 100),
    "width": 20,
    "counts": (counts := np.array([3, 13, 28, 14])),
    "coords": np.array(
        [
            [
                170 * (i + 1) + length * (1 / counts[:i] + 0.5 / counts[i]).sum(),
                650 - 170 / 2,
            ]
            for i in range(len(counts))
        ]
    ),
}
floors = [
    {
        "origin": (
            origin := np.array(
                [170 * i + (stairs["length"] / stairs["counts"][:i]).sum(), 0]
            )
        ),
        "areas": [
            {
                "origin": (suborigin := np.array([0, 0])),
                "wh": (wh := np.array([50, height])),
                "nodes": (nodes := generate_nodes(origin + suborigin, wh)),
                "start": (start := start + len(nodes)) - len(nodes),
            },
            {
                "origin": (suborigin := np.array([0, height])),
                "wh": (wh := np.array([50, 650 - height])),
                "nodes": (nodes := generate_nodes(origin + suborigin, wh)),
                "start": (start := start + len(nodes)) - len(nodes),
            },
            {
                "origin": (suborigin := np.array([50, height])),
                "wh": (wh := np.array([70, 650 - height])),
                "nodes": (nodes := generate_nodes(origin + suborigin, wh)),
                "start": (start := start + len(nodes)) - len(nodes),
            },
            {
                "origin": (suborigin := np.array([120, height])),
                "wh": (wh := np.array([50, 650 - height])),
                "nodes": (nodes := generate_nodes(origin + suborigin, wh)),
                "start": (start := start + len(nodes)) - len(nodes),
            },
            {
                "origin": (suborigin := np.array([120, 0])),
                "wh": (wh := np.array([50, height])),
                "nodes": (nodes := generate_nodes(origin + suborigin, wh)),
                "start": (start := start + len(nodes)) - len(nodes),
            },
        ],
    }
    for i, height in enumerate([400] + [480] * 4)
]
nodes = np.concatenate(
    [
        np.concatenate(tuple(map(lambda area: area["nodes"], floor["areas"])))
        for floor in floors
    ]
)
pd.DataFrame(nodes, columns=["x", "y"]).to_excel(
    os.path.join(root, "res", "结点坐标.xlsx"), index=False
)
stairs["start"] = people_nodes_count = len(nodes)
doors = {
    "start": stairs["start"] + len(stairs["coords"]),
    "nodes": [
        {
            "floor_index": (floor_index := 1),
            "area_index": 2,
            "coord": floors[floor_index]["origin"] + np.array([170 / 2, 480]),
        },
        {
            "floor_index": (floor_index := 2),
            "area_index": 4,
            "coord": floors[floor_index]["origin"] + np.array([50 + 70, 365]),
        },
        {
            "floor_index": (floor_index := 2),
            "area_index": 4,
            "coord": floors[floor_index]["origin"] + np.array([170, 365]),
        },
        {
            "floor_index": (floor_index := 3),
            "area_index": 0,
            "coord": floors[floor_index]["origin"] + np.array([0, 112]),
        },
    ],
}
nodes = np.concatenate((
    nodes, stairs["coords"], list(map(lambda node: node["coord"], doors["nodes"]))
))
danger = {
    "floor_index": (floor_index := 2),
    "area_index": (area_index := 2),
    "coord": (
        coord := floors[floor_index]["origin"]
        + floors[floor_index]["areas"][area_index]["origin"]
        + floors[floor_index]["areas"][area_index]["wh"] / 2
    ),
    "node_index": np.where(
        (coord == floors[floor_index]["areas"][area_index]["nodes"]).all(axis=1)
    )[0][0],
    "radius": 30,
}
# danger = {
#     "floor_index": (floor_index := np.random.randint(len(floors))),
#     "area_index": (area_index := np.random.randint(len(floors[floor_index]["areas"]))),
#     "node_index": (
#         node_index := np.random.randint(
#             len(floors[floor_index]["areas"][area_index]["nodes"])
#         )
#     ),
#     "coord": floors[floor_index]["areas"][area_index]["nodes"][node_index],
#     "radius": 30,
# }

fig = plt.figure()
ax = fig.add_subplot()
for floor in floors:
    for area in floor["areas"]:
        square = plt.Rectangle(
            floor["origin"] + area["origin"],
            area["wh"][0],
            area["wh"][1],
            fill=False,
        )
        ax.add_patch(square)
        x, y = zip(*area["nodes"])
        ax.scatter(x, y, s=1)
ax.set_xlim(0, 170 * 5 + stairs["length"] * (1 / stairs["counts"]).sum())
ax.set_ylim(0, 650)
fig.savefig(os.path.join(root, "res", "二维模型.pdf"))

graph = nx.Graph()

# 同一楼层同一区域的结点的边和相邻区域的结点的边
for floor_index, floor in enumerate(floors):
    for areai_index, areai in enumerate(floor["areas"]):
        for areaj_index, areaj in enumerate(
            floor["areas"][areai_index : areai_index + 2], areai_index
        ):
            for nodei_index, nodei in enumerate(areai["nodes"], areai["start"]):
                for nodej_index, nodej in enumerate(areaj["nodes"], areaj["start"]):
                    distance = np.sqrt(((nodei - nodej) ** 2).sum())
                    if floor_index == danger["floor_index"]:
                        i_to_danger = nodei - danger["coord"]
                        i_to_j = nodei - nodej
                        cos = np.dot(i_to_danger, i_to_j) / (np.linalg.norm(i_to_danger) * np.linalg.norm(i_to_j))
                        distance_to_danger = np.linalg.norm(i_to_danger) * np.sqrt(1 - cos ** 2)
                    else:
                        distance_to_danger = np.inf
                    if 0 < distance < viewpoint and distance_to_danger > danger["radius"]:
                        graph.add_edge(nodei_index, nodej_index, weight=distance)

# 楼梯的边
for (lfloor_index, lfloor), (rfloor_index, rfloor), (stair_index, stair) in zip(
    enumerate(floors[:-1]),
    enumerate(floors[1:], 1),
    enumerate(stairs["coords"], stairs["start"]),
):
    larea = lfloor["areas"][3]
    rarea = rfloor["areas"][1]
    for floor_index, area in zip([lfloor_index, rfloor_index], [larea, rarea]):
        min_distance = np.inf
        min_index = np.nan
        for node_index, node in enumerate(area["nodes"], area["start"]):
            distance_to_danger = (
                np.sqrt(((node - danger["coord"]) ** 2).sum())
                if floor_index == danger["floor_index"]
                else np.inf
            )
            ydistance = np.abs(stair[1] - node[1])
            distance = np.sqrt(((node - stair) ** 2).sum())
            if ydistance < stairs["width"] / 2 and distance < min_distance:
                min_distance = distance
                min_index = node_index
        if min_index != np.nan:
            graph.add_edge(stair_index, min_index, weight=distance)

# 门的边
for door_index, door in enumerate(doors["nodes"], doors["start"]):
    area = floors[door["floor_index"]]["areas"][door["area_index"]]
    for node_index, node in enumerate(area["nodes"], area["start"]):
        distance_to_danger = (
            np.sqrt(((node - danger["coord"]) ** 2).sum())
            if door["floor_index"] == danger["floor_index"]
            else np.inf
        )
        distance_to_door = np.sqrt(((node - door["coord"]) ** 2).sum())
        if distance_to_danger > danger["radius"] and distance_to_door < viewpoint:
            graph.add_edge(door_index, node_index, weight=distance_to_door)

pd.DataFrame(
    map(lambda edge: [*edge[:-1], edge[-1]["weight"]], graph.edges(data=True)),
    columns=["Node i", "Node j", "Weight"],
).to_excel(os.path.join(root, "res", "网络.xlsx"), index=False)

routes = []
for door_index, door in enumerate(doors["nodes"]):
    temp = nx.shortest_path(graph, target=door_index + doors["start"])
    numbers, pathes = list(zip(*sorted([(key, value) for key, value in temp.items() if key < len(nodes)], key=lambda x: x[0])))
    numbers = np.array(numbers)
    pathes = [nodes[path] for path in pathes]
    distances = [np.sqrt(((path[:-1] - path[1:]) ** 2).sum(axis=1)).sum() for path in pathes]
    route = pd.DataFrame(
        {
            "结点编号": numbers,
            "结点坐标": nodes[numbers].tolist(),
            "路径": pathes,
            "路径长度": distances,
        }
    )
    route = route[route["结点编号"] < people_nodes_count]
    fig = plt.figure()
    ax = fig.add_subplot()
    [ax.plot(path[:, 0], path[:, 1], linewidth=1) for path in route["路径"]]
    ax.scatter(
        door["coord"][0], door["coord"][1], label="door", color="black", zorder=2, s=5
    )
    # ax.scatter(danger["coord"][0], danger["coord"][1], label="danger", color="red", zorder=2, s=10)
    ax.add_artist(plt.Circle(danger["coord"], danger["radius"], label="danger", color="red", zorder=2))
    ax.legend(loc="lower right")
    ax.set_title(f"通过门{door_index}的路线")
    fig.savefig(os.path.join(root, "res", f"{ax.title.get_text()}.pdf"))
    routes.append(route)

distances = np.array([route["路径长度"] for route in routes])
solve = pd.DataFrame({
    "结点编号": routes[0]["结点编号"],
    "选择的门编号": (choosen := distances.argmin(axis=0)),
    "路径长度": (min_distances := distances.min(axis=0)),
})
result = pd.DataFrame(
    {
        "平均路径长度": [(s := min_distances.mean())],
        "耗费时间": [s / velocity_mean],
    }
)

with pd.ExcelWriter(os.path.join(root, "res", "最短路.xlsx")) as writer:
    pd.DataFrame({"number": range(len(doors["nodes"])), "coord": [node["coord"] for node in doors["nodes"]]}).to_excel(
        writer, sheet_name="door coordinate", index=False
    )
    [route.to_excel(writer, sheet_name=f"{i}", index=False) for i, route in enumerate(routes)]
    solve.to_excel(writer, sheet_name="最优化解", index=False)
    result.to_excel(writer, sheet_name="最优化结果", index=False)

fig = plt.figure()
ax = fig.add_subplot()
ax.add_artist(plt.Circle(danger["coord"], danger["radius"], label="danger", color="red", zorder=2))
[
    ax.scatter(door["coord"][0], door["coord"][1], label=f"door{door_index}", zorder=5, c=["red", "blue", "yellow", "green"][-door_index - 1], s=5)
    for door_index, door in enumerate(doors["nodes"])
]
ax.legend(loc="lower right")
[
    ax.plot(
        routes[door_index]["路径"].to_numpy()[node_index][:, 0],
        routes[door_index]["路径"].to_numpy()[node_index][:, 1],
        linewidth=1,
        color=["red", "blue", "yellow", "green"][door_index],
    )
    for node_index, door_index in enumerate(choosen)
]
ax.set_title("最优路径分配")
fig.savefig(os.path.join(root, "res", f"{ax.title.get_text()}.pdf"))

fig = plt.figure()
ax = fig.add_subplot()
ax.pie([(choosen == i).sum() for i in range(4)], labels=[f"door{i}" for i in range(4)], autopct='%1.1f%%')
ax.axis('equal')
ax.set_title("各门比例")
fig.savefig(os.path.join(root, "res", f"{ax.title.get_text()}.pdf"))
